'''创建noise_mask'''
import numpy as np
import cv2
from tqdm import tqdm

def events_to_count_images(event_list, img_shape=(240, 346), group_size=2000):
    """
    将事件流分组生成正负极性计数图
    :param event_list: 一维事件流列表 [p1, r1, c1, t1, p2, r2, c2, t2, ...]
    :param img_shape: 输出图像的形状 (H, W)
    :param group_size: 每组事件数
    :return: [(pos_img, neg_img), ...] 每组的正负极性计数图
    """
    results = []
    num_events = len(event_list) // 4
    for i in tqdm(range(0, num_events, group_size), desc="生成计数图"):
        pos_img = np.zeros(img_shape, dtype=np.int32)
        neg_img = np.zeros(img_shape, dtype=np.int32)
        end = min(i + group_size, num_events)
        for j in range(i, end):
            p = event_list[j*4]
            r = event_list[j*4+1]
            c = event_list[j*4+2]
            # t = event_list[j*4+3]  # 时间戳可忽略
            if 0 <= r < img_shape[0] and 0 <= c < img_shape[1]:
                if p == 1:
                    pos_img[r, c] += 1
                else:
                    neg_img[r, c] += 1
        results.append((pos_img, neg_img))
    return results

def events_to_count_image_single(event_list, img_shape=(240, 346)):
    """
    处理单组2000个事件, 生成正负极性计数图
    :param event_list: 一维事件流列表 [p1, r1, c1, t1, ...], 长度应为2000*4
    :param img_shape: 输出图像的形状 (H, W)
    :return: (pos_img, neg_img) 正负极性计数图
    """
    pos_img = np.zeros(img_shape, dtype=np.int32)
    neg_img = np.zeros(img_shape, dtype=np.int32)
    num_events = len(event_list) // 4
    for j in range(num_events):
        p = event_list[j*4]
        r = event_list[j*4+1]
        c = event_list[j*4+2]
        # t = event_list[j*4+3]  # 时间戳可忽略
        if 0 <= r < img_shape[0] and 0 <= c < img_shape[1]:
            if p == 1:
                pos_img[r, c] += 1
            else:
                neg_img[r, c] += 1
    return pos_img, neg_img

def create_noise_mask(pos_img, neg_img, kernel_size=7, thresh_noise=17):
    """
    根据正负极性计数图生成噪声掩码，并增加邻域无事件判定
    :param pos_img: 正极性事件计数图 (H, W)
    :param neg_img: 负极性事件计数图 (H, W)
    :param kernel_size: 滤波核大小
    :param thresh_noise: 噪声阈值
    :return: noise_mask (H, W), 噪声点为1, 其余为0
    """
    kernel = np.ones((kernel_size, kernel_size), dtype=np.float32)
    # 原有密度判定
    pos_density = cv2.filter2D(pos_img.astype(np.float32), -1, kernel, borderType=cv2.BORDER_REPLICATE)
    neg_density = cv2.filter2D(neg_img.astype(np.float32), -1, kernel, borderType=cv2.BORDER_REPLICATE)
    mask_density = ((pos_density < thresh_noise) & (neg_density < thresh_noise))
    # 新增：邻域无事件判定
    total_img = pos_img + neg_img
    neighbor_sum = cv2.filter2D((total_img > 0).astype(np.uint8), -1, kernel, borderType=cv2.BORDER_REPLICATE)
    mask_isolated = (neighbor_sum == 0)
    # 合并两种判定
    noise_mask = (mask_density | mask_isolated).astype(np.uint8)
    return noise_mask

def apply_noise_mask(pos_img, neg_img, noise_mask):
    """
    对正负极性计数图应用噪声掩码，噪声点置零
    :param pos_img: 正极性事件计数图
    :param neg_img: 负极性事件计数图
    :param noise_mask: 噪声掩码(噪声点为1, 其余为0)
    :return: 降噪后的正负极性计数图
    """
    pos_img_denoised = pos_img.copy()
    neg_img_denoised = neg_img.copy()
    pos_img_denoised[noise_mask == 1] = 0
    neg_img_denoised[noise_mask == 1] = 0
    return pos_img_denoised, neg_img_denoised
